markitect-main/markitect/finance/session_tracker.py

"""
Claude Session Cost Tracker for Issue Implementation.

This module tracks Claude session costs during issue implementation and generates
cost notes documenting the expenses associated with completing specific issues.
"""

import json
import os
from datetime import datetime, date
from decimal import Decimal
from typing import Dict, Any, Optional, List
from pathlib import Path

from .cost_manager import CostItemManager, CostItem
from .report_generator import CostReportGenerator, ReportConfig


class SessionCostTracker:
    """Tracks Claude session costs and generates cost notes for issues."""

    def __init__(self, db_path: str):
        """
        Initialize session cost tracker.

        Args:
            db_path: Path to SQLite database file
        """
        self.db_path = db_path
        self.cost_manager = CostItemManager(db_path)
        self.report_generator = CostReportGenerator(db_path)

    def estimate_session_cost(self, input_tokens: int, output_tokens: int,
                            model: str = "claude-sonnet-4") -> Dict[str, Any]:
        """
        Estimate Claude session cost based on token usage.

        Args:
            input_tokens: Number of input tokens used
            output_tokens: Number of output tokens generated
            model: Claude model used (affects pricing)

        Returns:
            Dictionary with cost breakdown
        """
        # Claude pricing (as of 2025 - these would be updated regularly)
        pricing = {
            "claude-sonnet-4": {
                "input_per_million": 3.00,  # $3 per million input tokens
                "output_per_million": 15.00  # $15 per million output tokens
            },
            "claude-sonnet-3.5": {
                "input_per_million": 3.00,
                "output_per_million": 15.00
            },
            "claude-haiku": {
                "input_per_million": 0.25,
                "output_per_million": 1.25
            }
        }

        if model not in pricing:
            model = "claude-sonnet-4"  # Default fallback

        rates = pricing[model]

        # Calculate costs in USD
        input_cost_usd = (input_tokens / 1_000_000) * rates["input_per_million"]
        output_cost_usd = (output_tokens / 1_000_000) * rates["output_per_million"]
        total_cost_usd = input_cost_usd + output_cost_usd

        # Convert to EUR (approximate rate - in practice would use real-time rates)
        usd_to_eur = 0.92  # Approximate conversion rate
        total_cost_eur = total_cost_usd * usd_to_eur

        return {
            "model": model,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": input_tokens + output_tokens,
            "input_cost_usd": round(input_cost_usd, 4),
            "output_cost_usd": round(output_cost_usd, 4),
            "total_cost_usd": round(total_cost_usd, 4),
            "total_cost_eur": round(total_cost_eur, 4),
            "pricing_rates": rates,
            "conversion_rate": usd_to_eur
        }

    def create_issue_cost_item(self, issue_id: int, issue_title: str,
                              session_cost: Dict[str, Any],
                              session_date: Optional[date] = None) -> Optional[int]:
        """
        Create a cost item for an issue implementation session.

        Args:
            issue_id: Issue ID number
            issue_title: Title of the issue
            session_cost: Cost breakdown from estimate_session_cost
            session_date: Date of the session (defaults to today)

        Returns:
            ID of created cost item, or None if creation failed
        """
        if session_date is None:
            session_date = date.today()

        # Get or create AI & ML Services category
        ai_category = self.cost_manager.get_category_by_name('AI & ML Services')
        if not ai_category:
            # Create the category if it doesn't exist
            category_id = self.cost_manager.create_category(
                'AI & ML Services',
                'Claude sessions and AI-powered development tools'
            )
        else:
            category_id = ai_category['id']

        # Create descriptive name and description
        name = f"Claude Session - Issue #{issue_id}"
        description = (
            f"Claude {session_cost['model']} session for implementing '{issue_title}'. "
            f"{session_cost['total_tokens']:,} tokens "
            f"({session_cost['input_tokens']:,} input, {session_cost['output_tokens']:,} output)"
        )

        # Create cost item
        cost_item = CostItem(
            category_id=category_id,
            name=name,
            description=description,
            cost_type='one_time',
            amount_eur=Decimal(str(session_cost['total_cost_eur'])),
            currency='EUR',
            starting_from_date=session_date
        )

        cost_item_id = self.cost_manager.create_cost_item(cost_item)

        if cost_item_id:
            print(f"✅ Created cost item for Issue #{issue_id}: €{session_cost['total_cost_eur']:.4f}")

        return cost_item_id

    def generate_issue_cost_note(self, issue_id: int, issue_title: str,
                               session_cost: Dict[str, Any],
                               implementation_summary: Optional[str] = None,
                               session_date: Optional[date] = None) -> str:
        """
        Generate a cost note document for an issue implementation.

        Args:
            issue_id: Issue ID number
            issue_title: Title of the issue
            session_cost: Cost breakdown from estimate_session_cost
            implementation_summary: Summary of what was implemented
            session_date: Date of the session

        Returns:
            Markdown content for the cost note
        """
        if session_date is None:
            session_date = date.today()

        # Prepare frontmatter
        frontmatter = {
            "note_type": "issue_cost_tracking",
            "issue_id": issue_id,
            "issue_title": issue_title,
            "session_date": session_date.isoformat(),
            "claude_model": session_cost['model'],
            "total_cost_eur": session_cost['total_cost_eur'],
            "total_cost_usd": session_cost['total_cost_usd'],
            "total_tokens": session_cost['total_tokens'],
            "generated_at": datetime.now().isoformat()
        }

        # Prepare contentmatter
        contentmatter = {
            "cost_tracking": {
                "issue": {
                    "id": issue_id,
                    "title": issue_title,
                    "implementation_date": session_date.isoformat()
                },
                "session": {
                    "model": session_cost['model'],
                    "token_usage": {
                        "input_tokens": session_cost['input_tokens'],
                        "output_tokens": session_cost['output_tokens'],
                        "total_tokens": session_cost['total_tokens']
                    },
                    "costs": {
                        "input_cost_usd": session_cost['input_cost_usd'],
                        "output_cost_usd": session_cost['output_cost_usd'],
                        "total_cost_usd": session_cost['total_cost_usd'],
                        "total_cost_eur": session_cost['total_cost_eur'],
                        "conversion_rate": session_cost['conversion_rate']
                    },
                    "pricing_rates": session_cost['pricing_rates']
                }
            }
        }

        # Build content
        content_lines = [
            f"# Issue #{issue_id} Implementation Cost",
            f"**Issue**: {issue_title}",
            f"**Date**: {session_date.strftime('%Y-%m-%d')}",
            f"**Claude Model**: {session_cost['model']}",
            "",
            "## Cost Summary",
            f"- **Total Cost**: €{session_cost['total_cost_eur']:.4f} (${session_cost['total_cost_usd']:.4f} USD)",
            f"- **Token Usage**: {session_cost['total_tokens']:,} tokens",
            f"- **Input Tokens**: {session_cost['input_tokens']:,} tokens @ ${session_cost['pricing_rates']['input_per_million']:.2f}/M",
            f"- **Output Tokens**: {session_cost['output_tokens']:,} tokens @ ${session_cost['pricing_rates']['output_per_million']:.2f}/M",
            "",
            "## Cost Breakdown",
            "",
            "| Component | Tokens | Rate ($/M) | Cost (USD) | Cost (EUR) |",
            "|-----------|--------|------------|------------|------------|",
            f"| Input | {session_cost['input_tokens']:,} | ${session_cost['pricing_rates']['input_per_million']:.2f} | ${session_cost['input_cost_usd']:.4f} | €{session_cost['input_cost_usd'] * session_cost['conversion_rate']:.4f} |",
            f"| Output | {session_cost['output_tokens']:,} | ${session_cost['pricing_rates']['output_per_million']:.2f} | ${session_cost['output_cost_usd']:.4f} | €{session_cost['output_cost_usd'] * session_cost['conversion_rate']:.4f} |",
            f"| **Total** | {session_cost['total_tokens']:,} | - | ${session_cost['total_cost_usd']:.4f} | €{session_cost['total_cost_eur']:.4f} |",
            "",
        ]

        if implementation_summary:
            content_lines.extend([
                "## Implementation Summary",
                implementation_summary,
                ""
            ])

        content_lines.extend([
            "## Cost Allocation",
            f"This cost has been allocated to the 'AI & ML Services' category as a one-time expense for issue #{issue_id} implementation.",
            "",
            "## Notes",
            f"- Currency conversion rate: 1 USD = {session_cost['conversion_rate']:.3f} EUR",
            f"- Pricing based on {session_cost['model']} rates as of {session_date}",
            "- Token counts and costs are estimates based on session usage",
        ])

        content = "\n".join(content_lines)

        return self._assemble_cost_note(frontmatter, content, contentmatter)

    def _assemble_cost_note(self, frontmatter: Dict[str, Any], content: str,
                           contentmatter: Dict[str, Any]) -> str:
        """Assemble complete cost note with frontmatter and contentmatter."""
        # Convert frontmatter to YAML
        fm_lines = ["---"]
        for key, value in frontmatter.items():
            if isinstance(value, str):
                fm_lines.append(f'{key}: "{value}"')
            else:
                fm_lines.append(f'{key}: {value}')
        fm_lines.append("---")

        # Convert contentmatter to HTML comment
        cm_json = json.dumps(contentmatter, indent=2)
        cm_lines = [
            "<!--",
            "contentmatter:",
            cm_json,
            "-->"
        ]

        # Assemble final document
        return "\n".join([
            "\n".join(fm_lines),
            "",
            content,
            "",
            "\n".join(cm_lines)
        ])

    def save_issue_cost_note(self, issue_id: int, cost_note_content: str,
                           output_dir: str = "cost_notes") -> str:
        """
        Save issue cost note to file.

        Args:
            issue_id: Issue ID number
            cost_note_content: Generated cost note content
            output_dir: Directory to save cost notes

        Returns:
            Path to saved file
        """
        # Create output directory
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)

        # Generate filename
        today = date.today().strftime('%Y-%m-%d')
        filename = f"issue_{issue_id:03d}_cost_{today}.md"
        file_path = output_path / filename

        # Save file
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(cost_note_content)

        return str(file_path)

    def track_issue_completion(self, issue_id: int, issue_title: str,
                             input_tokens: int, output_tokens: int,
                             model: str = "claude-sonnet-4",
                             implementation_summary: Optional[str] = None,
                             save_note: bool = True,
                             output_dir: str = "cost_notes") -> Dict[str, Any]:
        """
        Complete workflow: estimate cost, create cost item, generate note.

        Args:
            issue_id: Issue ID number
            issue_title: Title of the issue
            input_tokens: Number of input tokens used
            output_tokens: Number of output tokens generated
            model: Claude model used
            implementation_summary: Summary of implementation
            save_note: Whether to save the cost note to file
            output_dir: Directory to save cost notes

        Returns:
            Dictionary with tracking results
        """
        # Estimate session cost
        session_cost = self.estimate_session_cost(input_tokens, output_tokens, model)

        # Create cost item in database
        cost_item_id = self.create_issue_cost_item(issue_id, issue_title, session_cost)

        # Generate cost note
        cost_note = self.generate_issue_cost_note(
            issue_id, issue_title, session_cost, implementation_summary
        )

        # Save cost note if requested
        saved_path = None
        if save_note:
            saved_path = self.save_issue_cost_note(issue_id, cost_note, output_dir)

        return {
            "issue_id": issue_id,
            "issue_title": issue_title,
            "session_cost": session_cost,
            "cost_item_id": cost_item_id,
            "cost_note": cost_note,
            "saved_path": saved_path,
            "tracking_successful": cost_item_id is not None
        }

    def get_issue_costs_summary(self, issue_ids: Optional[List[int]] = None) -> Dict[str, Any]:
        """
        Get summary of costs for specific issues or all AI service costs.

        Args:
            issue_ids: Optional list of issue IDs to filter by

        Returns:
            Summary of issue implementation costs
        """
        # Get AI & ML Services category
        ai_category = self.cost_manager.get_category_by_name('AI & ML Services')
        if not ai_category:
            return {"total_costs": 0.0, "issue_count": 0, "items": []}

        # Get cost items in AI category
        items = self.cost_manager.list_cost_items(
            active_only=True,
            category_id=ai_category['id']
        )

        # Filter by issue IDs if specified
        if issue_ids:
            filtered_items = []
            for item in items:
                # Extract issue ID from name (format: "Claude Session - Issue #123")
                if "Issue #" in item['name']:
                    try:
                        item_issue_id = int(item['name'].split("Issue #")[1].split()[0])
                        if item_issue_id in issue_ids:
                            filtered_items.append(item)
                    except (IndexError, ValueError):
                        continue
            items = filtered_items

        total_cost = sum(float(item['amount_eur']) for item in items)

        return {
            "total_costs": total_cost,
            "issue_count": len(items),
            "items": items,
            "currency": "EUR"
        }